# 导入包
from matplotlib import pyplot as plt
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import statsmodels.api as sm
import numpy as np
import seaborn as sns

# 数据库配置
db_config = {
    'host': '127.0.0.1',
    'user': 'root',
    'password': 'root',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8mb4'   # 添加字符集设置
}
# 创建数据库连接引擎
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
# 使用pymysql连接数据库
conn = pymysql.connect(**db_config)

# 设置分块读取大小
chunk_size = 10000
# 查询指定股票代码和日期范围的数据
df = pd.read_sql_query("""
    SELECT d.* FROM date_1 d WHERE d.trade_date BETWEEN '2023-01-01' and '2023-12-31' and d.ts_code = '000001.SZ'
    """,
    conn,
    chunksize=chunk_size)

# 将分块读取的数据拼接为一个完整的DataFrame
df1 = pd.concat(df, ignore_index=True)

# 新增股票涨跌列，计算涨跌幅度
df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 4)  # 修正：小数点后保留4位

# 处理缺失值数据，删除'zd_closes'列中缺失值所在的行
df1 = df1.dropna(subset=['zd_closes'])

# 打印前几行数据
print(df1.head())

# 筛选自变量，排除不需要的列
ex = ['zd_closes', 'id', 'ts_code', 'trade_date', 'the_date', 'opens', 'high', 'low', 'closes', 'pre_closes', 'change']
# 获取数值型列的列名
number = df1.select_dtypes(include=['number']).columns.tolist()

# 从数值型列中移除不需要的列，得到自变量列名
newlist = [col for col in number if col not in ex]

# 构造回归公式
formulas = 'zd_closes ~ ' + ' + '.join(newlist)

# 拟合线性回归模型
res = smf.ols(formulas, data=df1).fit()

# 打印回归结果
print(res.summary())
#